import numpy as np
from scipy.linalg import pinv


def forward_kinematics(theta, DH_params):
    """
    计算机械臂的正向运动学。

    参数:
    theta (list): 当前关节角度列表。
    DH_params (numpy.array): DH参数矩阵，每行代表一个连杆的DH参数。

    返回:
    T (numpy.array): 末端执行器的变换矩阵。
    """
    T = np.eye(4)
    for i in range(len(theta)):
        alpha, a, d = DH_params[i, 1], DH_params[i, 2], DH_params[i, 3]
        alpha, theta_i = np.deg2rad(alpha), np.deg2rad(theta[i])
        A_i = np.array([
            [np.cos(theta_i), -np.sin(theta_i) * np.cos(alpha), np.sin(theta_i) * np.sin(alpha), a * np.cos(theta_i)],
            [np.sin(theta_i), np.cos(theta_i) * np.cos(alpha), -np.cos(theta_i) * np.sin(alpha), a * np.sin(theta_i)],
            [0, np.sin(alpha), np.cos(alpha), d],
            [0, 0, 0, 1]
        ])
        T = T @ A_i
    return T


def compute_F(T_current, T_des):
    """
    计算当前变换矩阵与目标变换矩阵之间的误差。

    参数:
    T_current (numpy.array): 当前变换矩阵。
    T_des (numpy.array): 目标变换矩阵。

    返回:
    F (numpy.array): 误差向量。
    """
    F = np.zeros(12)
    F[:3] = T_current[:3, 3] - T_des[:3, 3]
    F[3:6] = T_current[0, 2] - T_des[0, 2], T_current[1, 2] - T_des[1, 2], T_current[2, 2] - T_des[2, 2]
    F[6:9] = T_current[0, 1] - T_des[0, 1], T_current[1, 1] - T_des[1, 1], T_current[2, 1] - T_des[2, 1]
    F[9:] = T_current[0, 0] - T_des[0, 0], T_current[1, 0] - T_des[1, 0], T_current[2, 0] - T_des[2, 0]
    return F


def compute_jacobian(theta, DH_params, T_des):
    """
    计算雅可比矩阵。

    参数:
    theta (list): 当前关节角度列表。
    DH_params (numpy.array): DH参数矩阵。
    T_des (numpy.array): 目标变换矩阵。

    返回:
    J (numpy.array): 雅可比矩阵。
    """
    delta = 1e-6
    n = len(theta)
    m = 12
    J = np.zeros((m, n))
    T_current = forward_kinematics(theta, DH_params)
    F = compute_F(T_current, T_des)

    for i in range(n):
        theta_perturbed = theta.copy()
        theta_perturbed[i] += delta
        T_perturbed = forward_kinematics(theta_perturbed, DH_params)
        F_perturbed = compute_F(T_perturbed, T_des)
        J[:, i] = (F_perturbed - F) / delta

    return J


def inverse_kinematics(T_des, q0, DH_params, max_iter=100000, tol=1e-12):
    """
    使用牛顿迭代法求解逆运动学。

    参数:
    T_des (numpy.array): 目标变换矩阵。
    q0 (list): 初始关节角度列表。
    DH_params (numpy.array): DH参数矩阵。
    max_iter (int): 最大迭代次数。
    tol (float): 收敛容忍度。

    返回:
    q (numpy.array): 求解得到的关节角度列表。
    """
    q = q0
    for iter in range(max_iter):
        T = forward_kinematics(q, DH_params)
        F = compute_F(T, T_des)
        J = compute_jacobian(q, DH_params, T_des)
        J_pseudo_inv = pinv(J)
        q = q - J_pseudo_inv @ F
        q = np.mod(q + 90, 180) - 90  # 限制关节角在[-90, 90]度范围内
        if np.linalg.norm(F) < tol:
            break
    return q

if __name__ == '__main__':
# 示例使用
    DH_params = np.array([
        [1, -45, 0, 560, 10],
        [2, 45, 0, 450, 0],
        [3, 45, 0, 410, 10],
        [4, -45, 0, 380, -20],
        [5, 45, 0, 360, 0],
        [6, 0, 0, 210, 40]
    ])
    q0 = np.zeros(6)  # 初始关节角度向量
    T_des = np.eye(4)  # 目标变换矩阵，需要根据实际情况定义

    q = inverse_kinematics(T_des, q0, DH_params)
    print("牛顿法求得关节角向量为：")
    print(q)